#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
NumPy 高级用法示例
演示NumPy的高级功能和实用技巧
"""

import numpy as np
import time

def demonstrate_broadcasting():
    """演示广播机制"""
    print("=== 广播机制示例 ===")
    
    # 基本广播
    a = np.array([[1, 2, 3], [4, 5, 6]])  # (2, 3)
    b = np.array([10, 20, 30])             # (3,)
    
    print(f"数组a (2,3):\n{a}")
    print(f"数组b (3,): {b}")
    print(f"a + b (广播):\n{a + b}")
    
    # 更复杂的广播
    c = np.array([[1], [2], [3]])  # (3, 1)
    d = np.array([10, 20])         # (2,)
    
    print(f"\n数组c (3,1):\n{c}")
    print(f"数组d (2,): {d}")
    print(f"c + d (广播):\n{c + d}")
    
    # 实际应用：标准化数据
    data = np.random.randn(100, 5)
    normalized = (data - np.mean(data, axis=0)) / np.std(data, axis=0)
    print(f"\n数据标准化示例:")
    print(f"原数据形状: {data.shape}")
    print(f"均值: {np.mean(data, axis=0)}")
    print(f"标准化后均值: {np.mean(normalized, axis=0)}")

def demonstrate_fancy_indexing():
    """演示花式索引"""
    print("\n=== 花式索引示例 ===")
    
    arr = np.arange(20).reshape(4, 5)
    print(f"原数组 (4x5):\n{arr}")
    
    # 使用整数数组索引
    rows = np.array([0, 2, 3])
    cols = np.array([1, 3, 4])
    
    print(f"\n选择行 {rows}, 列 {cols}:")
    print(f"结果: {arr[rows, cols]}")
    
    # 选择特定行和所有列
    selected_rows = arr[[0, 2]]
    print(f"\n选择第0行和第2行:\n{selected_rows}")
    
    # 布尔索引与花式索引结合
    mask = arr > 10
    indices = np.where(mask)
    print(f"\n大于10的元素索引:")
    print(f"行索引: {indices[0]}")
    print(f"列索引: {indices[1]}")
    print(f"对应值: {arr[mask]}")

def demonstrate_structured_arrays():
    """演示结构化数组"""
    print("\n=== 结构化数组示例 ===")
    
    # 定义结构化数据类型
    student_dtype = np.dtype([
        ('name', 'U20'),      # Unicode字符串，最大20字符
        ('age', 'i4'),        # 32位整数
        ('height', 'f4'),     # 32位浮点数
        ('grade', 'f4')       # 32位浮点数
    ])
    
    # 创建结构化数组
    students = np.array([
        ('张三', 20, 175.5, 85.5),
        ('李四', 19, 168.0, 92.0),
        ('王五', 21, 180.2, 78.5),
        ('赵六', 20, 165.8, 88.0)
    ], dtype=student_dtype)
    
    print(f"学生数据:")
    print(students)
    
    # 访问特定字段
    print(f"\n所有学生姓名: {students['name']}")
    print(f"所有学生年龄: {students['age']}")
    print(f"平均身高: {np.mean(students['height']):.1f}cm")
    print(f"平均成绩: {np.mean(students['grade']):.1f}分")
    
    # 按成绩排序
    sorted_students = np.sort(students, order='grade')
    print(f"\n按成绩排序:")
    for student in sorted_students:
        print(f"{student['name']}: {student['grade']}分")

def demonstrate_matrix_operations():
    """演示矩阵运算"""
    print("\n=== 矩阵运算示例 ===")
    
    # 创建矩阵
    A = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    B = np.array([[9, 8, 7], [6, 5, 4], [3, 2, 1]])
    
    print(f"矩阵A:\n{A}")
    print(f"矩阵B:\n{B}")
    
    # 矩阵乘法
    C = np.dot(A, B)  # 或者 A @ B
    print(f"\n矩阵乘法 A @ B:\n{C}")
    
    # 转置
    print(f"\nA的转置:\n{A.T}")
    
    # 只对方阵有效的操作
    square_matrix = np.array([[2, 1], [1, 3]])
    print(f"\n方阵:\n{square_matrix}")
    
    # 行列式
    det = np.linalg.det(square_matrix)
    print(f"行列式: {det:.2f}")
    
    # 特征值和特征向量
    eigenvalues, eigenvectors = np.linalg.eig(square_matrix)
    print(f"特征值: {eigenvalues}")
    print(f"特征向量:\n{eigenvectors}")
    
    # 逆矩阵
    if det != 0:
        inverse = np.linalg.inv(square_matrix)
        print(f"逆矩阵:\n{inverse}")
        
        # 验证：矩阵乘以其逆矩阵应该等于单位矩阵
        identity_check = np.dot(square_matrix, inverse)
        print(f"验证 A * A^(-1):\n{identity_check}")

def demonstrate_polynomial_operations():
    """演示多项式操作"""
    print("\n=== 多项式操作示例 ===")
    
    # 定义多项式 p(x) = 2x^3 + 3x^2 - x + 5
    # NumPy中多项式系数从高次到低次排列
    coefficients = [2, 3, -1, 5]
    
    print(f"多项式系数 (从高次到低次): {coefficients}")
    print("对应多项式: 2x³ + 3x² - x + 5")
    
    # 计算多项式值
    x_values = np.linspace(-2, 2, 9)
    y_values = np.polyval(coefficients, x_values)
    
    print(f"\nx值: {x_values}")
    print(f"y值: {y_values}")
    
    # 多项式求根
    roots = np.roots(coefficients)
    print(f"\n多项式的根: {roots}")
    
    # 多项式拟合
    # 生成带噪声的数据点
    np.random.seed(42)
    x_data = np.linspace(-1, 1, 20)
    y_true = 2*x_data**2 + 3*x_data + 1
    y_noisy = y_true + 0.5 * np.random.randn(len(x_data))
    
    # 用2次多项式拟合
    fitted_coeffs = np.polyfit(x_data, y_noisy, 2)
    print(f"\n拟合得到的系数: {fitted_coeffs}")
    print(f"真实系数: [2, 3, 1]")

def demonstrate_fft_operations():
    """演示快速傅立叶变换"""
    print("\n=== 快速傅立叶变换示例 ===")
    
    # 创建信号：两个正弦波的叠加
    sampling_rate = 1000  # 采样率
    duration = 1.0        # 持续时间
    t = np.linspace(0, duration, int(sampling_rate * duration), endpoint=False)
    
    # 信号：50Hz + 120Hz + 噪声
    frequency1, frequency2 = 50, 120
    signal = (np.sin(2 * np.pi * frequency1 * t) + 
              0.5 * np.sin(2 * np.pi * frequency2 * t) + 
              0.2 * np.random.randn(len(t)))
    
    print(f"信号长度: {len(signal)}")
    print(f"采样率: {sampling_rate} Hz")
    print(f"包含频率: {frequency1} Hz, {frequency2} Hz")
    
    # 执行FFT
    fft_result = np.fft.fft(signal)
    frequencies = np.fft.fftfreq(len(signal), 1/sampling_rate)
    
    # 计算幅度谱
    magnitude = np.abs(fft_result)
    
    # 找到主要频率分量
    # 只看正频率部分
    positive_freq_indices = frequencies > 0
    positive_frequencies = frequencies[positive_freq_indices]
    positive_magnitudes = magnitude[positive_freq_indices]
    
    # 找到幅度最大的几个频率
    peak_indices = np.argsort(positive_magnitudes)[-5:]
    peak_frequencies = positive_frequencies[peak_indices]
    peak_magnitudes = positive_magnitudes[peak_indices]
    
    print(f"\n检测到的主要频率分量:")
    for freq, mag in zip(peak_frequencies, peak_magnitudes):
        if mag > np.max(positive_magnitudes) * 0.1:  # 只显示显著的峰值
            print(f"频率: {freq:.1f} Hz, 幅度: {mag:.2f}")

def demonstrate_random_sampling():
    """演示随机采样和分布"""
    print("\n=== 随机采样示例 ===")
    
    # 设置随机种子
    np.random.seed(42)
    
    # 各种分布
    print("不同概率分布的随机采样:")
    
    # 均匀分布
    uniform = np.random.uniform(0, 1, 5)
    print(f"均匀分布 [0,1]: {uniform}")
    
    # 正态分布
    normal = np.random.normal(0, 1, 5)
    print(f"标准正态分布: {normal}")
    
    # 指数分布
    exponential = np.random.exponential(2, 5)
    print(f"指数分布 (λ=0.5): {exponential}")
    
    # 泊松分布
    poisson = np.random.poisson(3, 10)
    print(f"泊松分布 (λ=3): {poisson}")
    
    # 二项分布
    binomial = np.random.binomial(10, 0.3, 10)
    print(f"二项分布 (n=10, p=0.3): {binomial}")
    
    # 随机选择
    choices = np.random.choice(['A', 'B', 'C', 'D'], size=10, p=[0.1, 0.2, 0.3, 0.4])
    print(f"随机选择: {choices}")
    
    # 随机排列
    arr = np.arange(10)
    np.random.shuffle(arr)
    print(f"随机排列: {arr}")

def demonstrate_memory_views():
    """演示内存视图和拷贝"""
    print("\n=== 内存视图和拷贝示例 ===")
    
    # 创建原始数组
    original = np.arange(20).reshape(4, 5)
    print(f"原始数组:\n{original}")
    
    # 视图（共享内存）
    view = original[1:3, 1:4]
    print(f"\n视图 (切片):\n{view}")
    
    # 修改视图会影响原数组
    view[0, 0] = 999
    print(f"修改视图后的原数组:\n{original}")
    
    # 深拷贝（独立内存）
    copy = original.copy()
    copy[0, 0] = 888
    print(f"\n修改拷贝后:")
    print(f"原数组:\n{original}")
    print(f"拷贝:\n{copy}")
    
    # 检查是否共享内存
    print(f"\n内存共享检查:")
    print(f"view与original共享内存: {np.shares_memory(view, original)}")
    print(f"copy与original共享内存: {np.shares_memory(copy, original)}")

def demonstrate_advanced_indexing():
    """演示高级索引技巧"""
    print("\n=== 高级索引技巧示例 ===")
    
    # 创建3D数组
    arr_3d = np.arange(24).reshape(2, 3, 4)
    print(f"3D数组 (2x3x4):\n{arr_3d}")
    
    # 使用省略号
    print(f"\n使用省略号选择所有第一维度:")
    print(f"arr_3d[..., 0]:\n{arr_3d[..., 0]}")
    
    # 使用np.where进行条件选择
    arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
    result = np.where(arr % 2 == 0, arr, -arr)
    print(f"\n条件选择示例:")
    print(f"原数组: {arr}")
    print(f"偶数保持，奇数变负: {result}")
    
    # 使用np.argmax和np.argmin
    matrix = np.random.randn(3, 4)
    print(f"\n随机矩阵:\n{matrix}")
    print(f"最大值位置: {np.unravel_index(np.argmax(matrix), matrix.shape)}")
    print(f"最小值位置: {np.unravel_index(np.argmin(matrix), matrix.shape)}")
    
    # 使用np.take进行索引
    indices = [0, 2, 1, 3]
    taken = np.take(arr, indices)
    print(f"\nusing np.take:")
    print(f"原数组: {arr}")
    print(f"索引 {indices}: {taken}")

def benchmark_operations():
    """性能基准测试"""
    print("\n=== 性能基准测试 ===")
    
    # 测试不同操作的性能
    size = 1000000
    a = np.random.randn(size)
    b = np.random.randn(size)
    
    # NumPy向量化操作
    start_time = time.time()
    result_numpy = a * b + np.sin(a)
    numpy_time = time.time() - start_time
    
    # Python循环操作（模拟）
    start_time = time.time()
    # 这里我们用NumPy的方式来模拟，因为真正的Python循环会太慢
    result_loop = np.array([a[i] * b[i] + np.sin(a[i]) for i in range(min(1000, size))])
    loop_time = (time.time() - start_time) * (size / min(1000, size))
    
    print(f"数组大小: {size}")
    print(f"NumPy向量化操作时间: {numpy_time:.6f}秒")
    print(f"Python循环操作预估时间: {loop_time:.6f}秒")
    print(f"性能提升倍数: {loop_time / numpy_time:.1f}x")

def main():
    """主函数，运行所有高级示例"""
    print("NumPy 高级功能演示")
    print("=" * 50)
    
    demonstrate_broadcasting()
    demonstrate_fancy_indexing()
    demonstrate_structured_arrays()
    demonstrate_matrix_operations()
    demonstrate_polynomial_operations()
    demonstrate_fft_operations()
    demonstrate_random_sampling()
    demonstrate_memory_views()
    demonstrate_advanced_indexing()
    benchmark_operations()

if __name__ == "__main__":
    main() 